package org.example.config;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ResultMap;
import org.apache.ibatis.mapping.ResultMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.example.util.PluginUtils;
import org.springframework.stereotype.Component;

import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * xxx
 *
 * @author lyl
 * @date 2023/12/31 16:09
 * @desc
 */
@Intercepts(
        value = {
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        }
)
@Slf4j(topic = "explainSql")
@Component
public class QueryExplainLogInterceptor implements Interceptor {
    protected static final Map<String, MappedStatement> countMsCache = new ConcurrentHashMap<>();
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object target = invocation.getTarget();
        Object[] args = invocation.getArgs();
        // 执行器
        final Executor executor = (Executor) target;
        // 解析参数
        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler resultHandler = (ResultHandler) args[3];
        // 获取boundSql
        BoundSql boundSql;
        if (args.length == 4) {
            boundSql = ms.getBoundSql(parameter);
        } else {
            // 几乎不可能走进这里面,除非使用Executor的代理对象调用query[args[6]]
            boundSql = (BoundSql) args[5];
        }

        toExplain(executor, ms, parameter, rowBounds, resultHandler, boundSql);

        return invocation.proceed();
    }

    /**
     *
     * @param executor
     * @param ms
     * @param parameter
     * @param rowBounds
     * @param resultHandler
     * @param boundSql
     * @return
     * @throws SQLException
     */
    public void toExplain(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds,
                          ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        MappedStatement explainMs = buildAutoExplainMappedStatement(ms);
        String explainSqlStr = autoExplainSql(boundSql.getSql());
        PluginUtils.MPBoundSql mpBoundSql = PluginUtils.mpBoundSql(boundSql);
        BoundSql explainSql = new BoundSql(explainMs.getConfiguration(), explainSqlStr, mpBoundSql.parameterMappings(), parameter);
        PluginUtils.setAdditionalParameter(explainSql, mpBoundSql.additionalParameters());

        CacheKey cacheKey = executor.createCacheKey(explainMs, parameter, rowBounds, explainSql);
        List<Object> resultList = executor.query(explainMs, parameter, rowBounds, resultHandler, cacheKey, explainSql);
        if (resultList != null && !resultList.isEmpty()) {
            for (Object o : resultList) {
                Map<Object, Object> map = (Map<Object, Object>) o;
                StringBuilder sb = new StringBuilder("explain info >>>>>>>> ");
                for (Map.Entry<Object, Object> entry : map.entrySet()) {
                    sb.append(entry.getKey()).append(": ").append(entry.getValue()).append(", ");
                }
                log.info(sb.append("<<<<<<<").toString());
            }

        }
    }

    /**
     * 构建自定义 MappedStatement
     *
     * @param ms MappedStatement
     * @return MappedStatement
     */
    protected MappedStatement buildAutoExplainMappedStatement(MappedStatement ms) {
        final String explainId = ms.getId() + "_cusExplain";
        final Configuration configuration = ms.getConfiguration();
        return computeIfAbsent(countMsCache, explainId, key -> {
            MappedStatement.Builder builder = new MappedStatement.Builder(configuration, key, ms.getSqlSource(), ms.getSqlCommandType());
            builder.resource(ms.getResource());
            builder.fetchSize(ms.getFetchSize());
            builder.statementType(ms.getStatementType());
            builder.timeout(ms.getTimeout());
            builder.parameterMap(ms.getParameterMap());
            builder.resultMaps(Collections.singletonList(new ResultMap.Builder(configuration, "cusExplain", Map.class, Collections.emptyList()).build()));
            builder.resultSetType(ms.getResultSetType());
            builder.cache(ms.getCache());
            builder.flushCacheRequired(ms.isFlushCacheRequired());
            builder.useCache(ms.isUseCache());
            return builder.build();
        });
    }

    private String autoExplainSql(String sql) {
        return "explain (" + sql + ")";
    }

    private static <K, V> V computeIfAbsent(Map<K, V> concurrentHashMap, K key, Function<? super K, ? extends V> mappingFunction) {
        V v = concurrentHashMap.get(key);
        if (v != null) {
            return v;
        }
        return concurrentHashMap.computeIfAbsent(key, mappingFunction);
    }


}
